package main

import (
	"bufio"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"runtime"
	"strconv"
	"sync"
	"time"
)

func main() {
	// 解析命令行参数
	var totalRecords int
	var outputPath string
	flag.IntVar(&totalRecords, "n", 1_000_000_000, "要生成的记录数量")
	flag.StringVar(&outputPath, "o", "data.txt", "输出文件路径")
	flag.Parse()

	if totalRecords <= 0 {
		fmt.Println("错误: 记录数量必须大于0")
		os.Exit(1)
	}

	fmt.Printf("开始生成 %d 条记录到: %s\n", totalRecords, outputPath)
	start := time.Now()

	// 创建输出文件
	f, err := os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		fmt.Printf("创建文件失败: %v\n", err)
		os.Exit(1)
	}
	defer f.Close()

	// 使用带缓冲的写入器
	writer := bufio.NewWriterSize(f, 64*1024*1024) // 64MB缓冲区
	defer writer.Flush()

	// 计算任务数量
	chunkSize := 10_000_000                                 // 每个任务处理1000万条
	numChunks := (totalRecords + chunkSize - 1) / chunkSize // 向上取整

	// 创建任务通道
	taskChan := make(chan int, numChunks)
	for i := 0; i < numChunks; i++ {
		taskChan <- i
	}
	close(taskChan)

	// 启动工作池
	var wg sync.WaitGroup
	workers := runtime.NumCPU() * 2 // CPU核心数×2
	fmt.Printf("使用 %d 个工作线程\n", workers)

	for w := 0; w < workers; w++ {
		wg.Add(1)
		go workerv1(&wg, taskChan, writer, totalRecords, chunkSize)
	}

	// 等待所有任务完成
	wg.Wait()

	duration := time.Since(start)
	fmt.Printf("生成完成! 耗时: %v\n", duration)
	fmt.Printf("平均速度: %.2f 万条/秒\n",
		float64(totalRecords)/duration.Seconds()/10000)
}

func workerv1(wg *sync.WaitGroup, tasks <-chan int, writer *bufio.Writer, totalRecords, chunkSize int) {
	defer wg.Done()

	// 每个worker维护独立随机源
	rng := rand.New(rand.NewSource(time.Now().UnixNano()))

	// 线程本地缓冲区
	buf := make([]byte, 0, 1024*1024) // 预分配1MB

	for chunkIndex := range tasks {
		startID := chunkIndex*chunkSize + 1
		endID := startID + chunkSize - 1
		if endID > totalRecords {
			endID = totalRecords
		}
		chunkRecords := endID - startID + 1

		for i := 0; i < chunkRecords; i++ {
			id := startID + i

			// 直接生成格式化的字节序列
			buf = append(buf, strconv.Itoa(id)...) // ID
			buf = append(buf, ',')

			c1 := rng.Intn(10000) // 0-9999随机整数
			buf = append(buf, strconv.Itoa(c1)...)
			buf = append(buf, ',')

			c2 := float64(rng.Intn(100000)) / 100.0 // 0.00-999.99随机浮点
			c2Str := strconv.FormatFloat(c2, 'f', 2, 64)
			buf = append(buf, c2Str...)
			buf = append(buf, '\n')

			// 批量写入
			if len(buf) >= cap(buf)-128 {
				_, err := writer.Write(buf)
				if err != nil {
					fmt.Printf("写入错误: %v\n", err)
				}
				buf = buf[:0] // 重置缓冲区
			}
		}

		// 写剩余数据
		if len(buf) > 0 {
			_, err := writer.Write(buf)
			if err != nil {
				fmt.Printf("写入错误: %v\n", err)
			}
			buf = buf[:0]
		}

		// 进度报告
		if chunkIndex%10 == 0 {
			fmt.Printf("已完成区块: %d/%d (ID: %d-%d)\n",
				chunkIndex+1, cap(tasks), startID, endID)
		}
	}
}
